[logging] Split rollout trajectory time into inference-engine wait vs env.step#1900
Open
eicherseiji wants to merge 14 commits into
Open
[logging] Split rollout trajectory time into inference-engine wait vs env.step#1900eicherseiji wants to merge 14 commits into
eicherseiji wants to merge 14 commits into
Conversation
… env.step Rollout timing is currently reported only as a total per trajectory (`generate/trajectory_completion_time_*`, added in NovaSky-AI#1804). That total cannot distinguish an engine-bound rollout from an environment-bound one, which is the first question to ask when GPU utilization sags during generation. Accumulate, per trajectory in `agent_loop`, the wall-clock time spent awaiting `inference_engine_client.generate()` and the time spent in `env.step()`, and report: generate/trajectory_llm_time_{mean,p90,max} generate/trajectory_env_time_{mean,p90,max} generate/frac_time_in_env `frac_time_in_env` is time-weighted (sum over sum) so long trajectories count proportionally rather than each trajectory contributing equally. The two components need not sum to `e2e_time`; the remainder is tokenization, chat-template rendering and other in-loop bookkeeping. Both are plumbed like the existing `e2e_time`: optional, and omitted entirely if any trajectory in the batch did not record them. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… gauge, buffer depth Three observability additions to the fully-async RL trainer, all motivated by a post-mortem where reconstructing GPU idle time required correlating Prometheus (GPU/disk, wall-clock keyed) against W&B (phase timings, step keyed) by hand. 1. Wire RayGpuMonitor into the async loop. It is constructed in the base RayPPOTrainer and flushed in the base loop, but FullyAsyncRayPPOTrainer overrides train() and never started or flushed it -- so runs on the async trainer logged zero `ray/` GPU keys despite `enable_ray_gpu_monitor=True`. Start it at loop entry, flush per step into the committed payload, stop it in the finally. 2. Publish a `skyrl_training_phase` gauge via `ray.util.metrics` (skyrl/train/utils/phase_metrics.py). Ray exports these to the same Prometheus that scrapes node GPU metrics, so `avg(ray_node_gpus_utilization) by (Phase)` becomes a single-store query instead of a manual cross-tracker correlation -- and it works after a cluster restart, when the experiment tracker may be unreachable. Set at each existing Timer() phase boundary (waiting_for_buffer / converting / training / weight_sync / eval / checkpoint); best-effort, never raises. 3. Log `async/gen_buffer_qsize` and `async/gen_buffer_maxsize` at step start. Run-ahead depth was previously only in a tqdm postfix; as a metric it distinguishes a throughput-limited generator (buffer low) from a gate-limited one (buffer at the staleness-bounded maxsize). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Item 1's llm/env time split was computed only at the per-generate()
get_rollout_metrics call. Every logging path (async trainer, eval,
step-wise) re-aggregates through concatenate_generator_outputs, which had
no raw llm/env lists and fell back to the substring-based extra_keys
aggregator. There p90 and frac_time_in_env hit the sum() branch, so a
16-group mini-batch logged frac_time_in_env near 12 and p90 near 16x its
real value.
Store the raw per-trajectory llm/env lists on GeneratorOutput like
e2e_time, thread them through concatenate_generator_outputs with the same
step-wise last-step filter, and recompute all seven metrics there so the
extra_keys fallback leaves them alone.
Also in this commit:
- Extend test_llm_vs_env_time_split_metrics to assert the split on a
concatenated output, the path that actually broke.
- Fix the phase_metrics module docstring query. The gauge has a lowercase
`phase` tag and no shared label with the GPU metric, so
`avg(ray_node_gpus_utilization) by (Phase)` does not run; replace it
with a vector match against `ray_skyrl_training_phase{phase="eval"}`.
- Deduplicate the last-step filter and the frac-block array construction.
Add ScalarGauges, a best-effort ray.util.metrics scalar-gauge helper that
no-ops without Ray, like TrainingPhaseGauge. Publish the async run-ahead
buffer levels through it so buffer pressure sits in the same Prometheus
store as the training-phase gauge and node GPU metrics, not only in the
experiment tracker:
skyrl_gen_buffer_qsize completed groups buffered at step start
skyrl_gen_buffer_maxsize staleness-bounded buffer capacity
skyrl_mini_batch_size groups consumed per training step
skyrl_gen_group_keep_rate kept / (kept + dropped) for the last
mini-batch, the zero-variance drop rate
under sample_full_batch
Rename the run-ahead phrasing to plain buffer depth in the buffer-metric and gauge comments; no behavior change.
The async-trainer observability items bundled here (RayGpuMonitor wiring, training-phase gauge, buffer depth) move to their own PRs so each is a reviewable unit. This restores fully_async_trainer.py to base and removes phase_metrics.py and its tests. The rollout LLM vs env time split and its aggregation across concatenate_generator_outputs stay.
This was referenced Jul 17, 2026
SumanthRH
pushed a commit
that referenced
this pull request
Jul 17, 2026
) ## What does this PR do? Fixes a wiring bug: `RayGpuMonitor` never ran on the fully async trainer. `RayGpuMonitor` (#1712) is constructed in `RayPPOTrainer.__init__` and enabled by default with `trainer.enable_ray_gpu_monitor=True`. The base loop starts it, flushes it into the per-step log payload, and stops it. `FullyAsyncRayPPOTrainer` overrides `train()` and never called any of the three, so fully async runs logged no `ray/` GPU keys at all. Confirmed on real runs. The async loop now starts the monitor at loop entry, flushes it into the per-step timing payload, and stops it in the `finally`, mirroring the base loop wiring. ## Tests Wiring only, covered by lint and existing tests. Exercising it end to end needs a GPU cluster. A structural lifecycle test that catches any base-loop hook dropped by an overriding trainer is a planned follow-up. Split out of #1900.
Review feedback pass: - Add trajectory_llm_times/env_times to the GeneratorOutput field guardrail test, which this branch had broken. - Extract _add_time_stats for the mean/p90/max blocks (completion, llm, env) and reuse its sums for frac_time_in_env. - Extract _optional_times for the omit-if-any-None per-prompt lists. - Use _concat_field for stop_reasons and rollout_logprobs too. - Cut duplicated comments; the llm/env split was explained in six places, now documented once per surface. - Assert the step-wise per-step replication of the split in the existing step-wise timing test.
Review feedback: the llm and env times traveled as parallel named fields (llm_time/env_time on two dataclasses, trajectory_llm_times/ trajectory_env_times on GeneratorOutput), so every hop re-documented and re-plumbed each field, and the omit-if-any-None handling needed its own helper per field. One time_splits dict per trajectory (keys "llm", "env") and one trajectory_time_splits dict-of-lists on GeneratorOutput replace them. The split semantics are documented once, on TrajectoryOutput.time_splits. _optional_times is gone; get_rollout_metrics loops the dict through _add_time_stats generically, so a future component is a new key rather than a new field at every layer. Metric names are unchanged.
frac_time_in_env is derivable exactly as env_mean / (llm_mean + env_mean)
from metrics already logged, so it added convenience, not information.
Replace the "need not sum to e2e" caveat with the metric that closes the
stack: generate/trajectory_overhead_time_{mean,p90,max}, computed as e2e
minus every recorded split component. Generic over the split keys, so a
future component reallocates time out of overhead without touching this
code. The test asserts overhead is the exact remainder and survives
concatenation.
The generate() harness test conflated timing attribution with metric math. Attribution keeps the integration test, now reading the splits straight off GeneratorOutput instead of spying on get_rollout_metrics. The metric math and concatenation move to exact-value tests over hand-built outputs in test_generator_output_utils.py. Also drop the overhead code comment.
The dataclass fields carry no comment; base.py states the keys and that the whole field is None (never partially populated) when any trajectory did not record its split.
eicherseiji
marked this pull request as ready for review
July 18, 2026 00:26
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces tracking and aggregation of trajectory generation time splits (e.g., distinguishing between LLM and environment step times) to compute detailed completion and overhead metrics. The review feedback highlights three potential runtime issues in the newly added utility functions: a potential KeyError in _split_lists due to mismatched keys, a potential TypeError in _concat_field if subsequent outputs contain None values, and a potential TypeError in _concat_time_splits if any trajectory lacks time splits. Robust code suggestions were provided to safely handle these edge cases.
_concat_time_splits keyed off the first output, so concatenating a batch that recorded splits with one that did not crashed when the missing batch came second. Fold list and dict-of-lists handling into one _concat_optional_field with any-missing-is-None semantics for every optional field, and let _last_step_only recurse into dicts. Adds an order-independence regression test.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Splits each trajectory's rollout time into engine time and env time. Today only the total exists (
trajectory_completion_time_*, #1804), so a slow rollout could be engine-bound or env-bound and the metrics cannot tell which.New metrics:
generate/trajectory_llm_time_{mean,p90,max}: time in the inference-engine call, summed over turnsgenerate/trajectory_env_time_{mean,p90,max}: time inenv.step(), summed over turnsgenerate/trajectory_overhead_time_{mean,p90,max}: everything else ine2e_time(env construction and teardown, tokenization, chat templating, event-loop scheduling), so the three bands sum to the total and unattributed time is visibleHow
agent_loopaccumulates onetime_splitsdict per trajectory (keysllm,env). It travels as a singletrajectory_time_splitsdict-of-lists onGeneratorOutput. Stats are recomputed from the raw lists inconcatenate_generator_outputs, which every logging path goes through, because per-group aggregates like a p90 cannot be combined. Overhead is e2e minus every recorded split component.Scope
#1925 (env setup time band) is stacked on this PR.
Tests
test_llm_vs_env_time_split_metrics(new): runsgenerate()with an env deliberately slower than the engine and asserts each sleep lands in the right component.test_generator_output_utils.py. Step-wise replication and theGeneratorOutputfield guardrail are asserted in existing tests.uv run --isolated --extra dev --extra skyrl-train pytest tests/train/generators/passes, 51 tests. ruff and black clean on changed files.